"""
ValueMomentumComposite Strategy Implementation for Zipline
Combines value factors (PE, PB, PS) with momentum factor
"""

from zipline.api import order_target_percent, symbol, get_datetime
from zipline.finance import commission, slippage
import pandas as pd

def initialize(context):
    # Strategy parameters
    context.PE_threshold = 15
    context.PB_threshold = 2
    context.PS_threshold = 1.5
    context.momentum_window = 252  # 12 months trading days
    context.momentum_threshold = 0.3  # Top 30%
    context.max_positions = 25
    context.max_position_size = 0.05  # 5% per stock
    context.stop_loss_daily = 0.02  # 2% daily stop loss
    context.stop_loss_monthly = 0.08  # 8% monthly stop loss
    
    # Rebalancing tracking
    context.last_rebalance = None
    context.last_month = None
    
    # Set commission and slippage
    context.set_commission(commission.PerShare(cost=0.005, min_trade_cost=1))
    context.set_slippage(slippage.FixedSlippage(spread=0.005))

def handle_data(context, data):
    # Check if we need to rebalance (quarterly)
    current_date = get_datetime().normalize()
    
    # Initialize last_rebalance if not set
    if context.last_rebalance is None:
        context.last_rebalance = current_date
    
    # Quarterly rebalancing (every 3 months)
    if (current_date - context.last_rebalance).days >= 90:
        rebalance(context, data)
        context.last_rebalance = current_date
    
    # Monthly return tracking
    current_month = current_date.month
    if context.last_month != current_month:
        context.portfolio.starting_value = context.portfolio.portfolio_value
        context.last_month = current_month
    
    # Daily stop loss check
    stop_loss_check(context, data)

def rebalance(context, data):
    # Get fundamental data and filter stocks
    universe = get_fundamental_data(context, data)
    momentum = calculate_momentum(context, data, universe)
    selected = filter_stocks(context, universe, momentum)
    
    # Exit positions not in selected stocks
    for stock in context.portfolio.positions:
        if stock not in selected:
            order_target_percent(stock, 0)
    
    # Allocate equal weights to selected stocks
    weight = min(context.max_position_size, 1.0 / len(selected))
    for stock in selected:
        order_target_percent(stock, weight)

def get_fundamental_data(context, data):
    """Retrieve stocks meeting fundamental criteria"""
    # In production, replace with actual fundamental data source
    # This is a placeholder implementation
    universe = [
        stock for stock in data.fetcher_assets
        if data.current(stock, 'pe_ratio') < context.PE_threshold
        and data.current(stock, 'pb_ratio') < context.PB_threshold
        and data.current(stock, 'ps_ratio') < context.PS_threshold
    ]
    return universe

def calculate_momentum(context, data, universe):
    """Calculate 12-month momentum for stocks"""
    momentum = {}
    for stock in universe:
        # Calculate 12-month return
        prices = data.history(stock, 'price', context.momentum_window, '1d')
        momentum[stock] = (prices[-1] / prices[0]) - 1
    return momentum

def filter_stocks(context, universe, momentum):
    """Filter stocks based on momentum threshold"""
    # Sort by momentum and select top 30%
    sorted_stocks = sorted(universe, key=lambda x: momentum[x], reverse=True)
    top_count = max(1, int(len(sorted_stocks) * context.momentum_threshold))
    return sorted_stocks[:min(top_count, context.max_positions)]

def stop_loss_check(context, data):
    """Implement daily and monthly stop loss logic"""
    portfolio = context.portfolio
    
    # Calculate daily returns
    daily_return = (portfolio.returns[-1] if len(portfolio.returns) > 0 else 0)
    
    # Calculate monthly returns
    monthly_return = (portfolio.portfolio_value / context.portfolio.starting_cash) - 1
    
    # Implement stop loss rules
    if daily_return < -context.stop_loss_daily:
        reduce_risk_exposure(context, data)
    if monthly_return < -context.stop_loss_monthly:
        liquidate_positions(context, data)

def reduce_risk_exposure(context, data):
    """Reduce position sizes by 50%"""
    for stock in context.portfolio.positions:
        current_position = context.portfolio.positions[stock].amount
        target_position = current_position * 0.5
        order_target_percent(stock, target_position)

def liquidate_positions(context, data):
    """Liquidate all positions"""
    for stock in context.portfolio.positions:
        order_target_percent(stock, 0)